#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
彩票号码统计分析工具

功能说明：
1. 出现次数(occurrence_count): 统计每个号码在历史开奖中出现的总次数
2. 当前遗漏(current_omission): 当前连续未出现的期数
3. 上期遗漏(last_omission): 上一期的遗漏值
4. 最大遗漏(max_omission): 历史以来最大的遗漏值
5. 平均遗漏(avg_omission): 历史平均遗漏值 = 总期数 / 出现次数
6. 最大连出(max_consecutive): 连续出现的最大次数

使用方法：
    # 计算双色球统计
    python number_statistics_calculator.py --lottery ssq
    
    # 计算大乐透统计并强制重算
    python number_statistics_calculator.py --lottery dlt --recalculate
    
    # 计算所有彩种
    python number_statistics_calculator.py --all
"""

import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from crawler.database import db_manager
from crawler.config.settings import LOTTERY_TYPES
from loguru import logger
from typing import Dict, List, Set, Tuple
from collections import defaultdict
import argparse
import json


class NumberStatisticsCalculator:
    """号码统计计算器 - 高效批量计算算法"""
    
    def __init__(self, lottery_type: str):
        self.lottery_type = lottery_type
        self.lottery_config = LOTTERY_TYPES.get(lottery_type, {})
        self.lottery_name = self.lottery_config.get('name', lottery_type)
        
    def get_number_ranges(self) -> Dict[str, range]:
        """获取号码范围（不同彩种有不同的号码范围）"""
        ranges = {}
        
        if self.lottery_type == 'ssq':
            ranges['red'] = range(1, 34)    # 红球 1-33
            ranges['blue'] = range(1, 17)   # 蓝球 1-16
        elif self.lottery_type == 'dlt':
            ranges['red'] = range(1, 36)    # 前区 1-35
            ranges['blue'] = range(1, 13)   # 后区 1-12
        elif self.lottery_type == 'fc3d':
            ranges['all'] = range(0, 10)    # 0-9
        elif self.lottery_type == 'pl3':
            ranges['all'] = range(0, 10)    # 0-9
        elif self.lottery_type == 'pl5':
            ranges['all'] = range(0, 10)    # 0-9
        elif self.lottery_type == 'qlc':
            ranges['red'] = range(1, 31)    # 基本号 1-30
            ranges['blue'] = range(1, 31)   # 特别号 1-30
        elif self.lottery_type == 'qxc':
            ranges['all'] = range(0, 10)    # 0-9
        elif self.lottery_type == 'kl8':
            ranges['all'] = range(1, 81)    # 1-80
        else:
            logger.error(f"不支持的彩种: {self.lottery_type}")
            return {}
        
        return ranges
    
    def fetch_historical_data(self) -> List[Dict]:
        """获取历史开奖数据（按期号正序排列）"""
        sql = """
        SELECT period, draw_date, red_numbers, blue_numbers, all_numbers
        FROM lottery_results
        WHERE lottery_type = %s
        ORDER BY period ASC
        """
        
        results = db_manager.execute_query(sql, (self.lottery_type,))
        logger.info(f"获取到 {len(results)} 期历史数据")
        return results
    
    def parse_numbers(self, record: Dict) -> Tuple[Set[int], Set[int], Set[int]]:
        """解析一期的号码数据
        
        Returns:
            (red_numbers, blue_numbers, all_numbers)
        """
        red_numbers = set()
        blue_numbers = set()
        all_numbers = set()
        
        # 优先使用red_numbers和blue_numbers字段
        if record.get('red_numbers'):
            try:
                red_list = record['red_numbers'] if isinstance(record['red_numbers'], list) else json.loads(record['red_numbers'])
                red_numbers = set(int(n) for n in red_list)
            except:
                pass
        
        if record.get('blue_numbers'):
            try:
                blue_list = record['blue_numbers'] if isinstance(record['blue_numbers'], list) else json.loads(record['blue_numbers'])
                blue_numbers = set(int(n) for n in blue_list)
            except:
                pass
        
        # 如果red_numbers和blue_numbers为空，尝试从all_numbers解析
        if not red_numbers and not blue_numbers and record.get('all_numbers'):
            try:
                all_data = record['all_numbers']
                if isinstance(all_data, str):
                    all_data = json.loads(all_data)
                
                if isinstance(all_data, list):
                    # 双色球/大乐透格式：["红球", "蓝球"]
                    if len(all_data) >= 2:
                        # 解析红球（第一部分，空格或逗号分隔）
                        red_str = str(all_data[0])
                        red_numbers = set(int(n.strip()) for n in red_str.replace(',', ' ').split() if n.strip().isdigit())
                        
                        # 解析蓝球（第二部分，空格或逗号分隔）
                        blue_str = str(all_data[1])
                        blue_numbers = set(int(n.strip()) for n in blue_str.replace(',', ' ').split() if n.strip().isdigit())
                    
                    # 单区彩种格式：直接是号码数组
                    elif len(all_data) >= 1 and self.lottery_type in ['fc3d', 'pl3', 'pl5', 'qxc']:
                        # 可能是空格分隔的字符串或数组
                        if isinstance(all_data[0], str):
                            all_numbers = set(int(n.strip()) for n in str(all_data[0]).replace(',', ' ').split() if n.strip().isdigit())
                        else:
                            all_numbers = set(int(n) for n in all_data)
            except Exception as e:
                pass
        
        # 如果还没有解析到any data，尝试直接解析all_numbers为单区号码
        if not red_numbers and not blue_numbers and not all_numbers and record.get('all_numbers'):
            try:
                all_data = record['all_numbers']
                if isinstance(all_data, str):
                    all_data = json.loads(all_data)
                if isinstance(all_data, list):
                    all_numbers = set(int(n) for n in all_data if str(n).strip().isdigit())
            except:
                pass
        
        return red_numbers, blue_numbers, all_numbers
    
    def calculate_statistics(self, historical_data: List[Dict]) -> Dict:
        """计算所有统计指标 - 使用高效的单遍扫描算法
        
        算法复杂度：O(n * m)，其中n是期数，m是号码数量
        空间复杂度：O(m)
        
        返回格式：
        {
            'red': {
                1: {
                    'number': 1,
                    'occurrence_count': 100,
                    'current_omission': 5,
                    'last_omission': 3,
                    'max_omission': 25,
                    'avg_omission': 2.5,
                    'max_consecutive': 3
                },
                ...
            },
            'blue': {...}
        }
        """
        if not historical_data:
            logger.warning("没有历史数据可供分析")
            return {}
        
        # 获取号码范围
        ranges = self.get_number_ranges()
        if not ranges:
            return {}
        
        logger.info(f"开始计算统计数据，共 {len(historical_data)} 期")
        
        # 初始化统计数据结构
        stats = {}
        for zone, number_range in ranges.items():
            stats[zone] = {}
            for number in number_range:
                stats[zone][number] = {
                    'number': number,
                    'occurrence_count': 0,
                    'current_omission': 0,
                    'last_omission': 0,
                    'max_omission': 0,
                    'avg_omission': 0.0,
                    'max_consecutive': 0,
                    # 辅助计算字段（计算完成后会删除）
                    '_current_consecutive': 0,    # 当前连出次数
                    '_last_appeared_period': -1,  # 最后出现的期数索引
                }
        
        total_periods = len(historical_data)
        
        # 单遍扫描算法：遍历每一期数据
        for period_index, record in enumerate(historical_data):
            period = record['period']
            
            # 解析本期号码
            red_numbers, blue_numbers, all_numbers = self.parse_numbers(record)
            
            # 每1000期输出一次进度
            if (period_index + 1) % 1000 == 0:
                logger.info(f"处理进度: {period_index + 1}/{total_periods}")
            
            # 更新每个区域的统计数据
            for zone, number_range in ranges.items():
                # 确定当前区域的开奖号码
                if zone == 'red':
                    drawn_numbers = red_numbers
                elif zone == 'blue':
                    drawn_numbers = blue_numbers
                else:  # 'all'
                    drawn_numbers = all_numbers
                
                # 遍历该区域的所有号码
                for number in number_range:
                    stat = stats[zone][number]
                    
                    if number in drawn_numbers:
                        # ===== 号码出现 =====
                        
                        # 1. 增加出现次数
                        stat['occurrence_count'] += 1
                        
                        # 2. 连出计数+1
                        stat['_current_consecutive'] += 1
                        
                        # 3. 更新最大连出
                        if stat['_current_consecutive'] > stat['max_consecutive']:
                            stat['max_consecutive'] = stat['_current_consecutive']
                        
                        # 4. 计算本次遗漏值（如果之前出现过）
                        if stat['_last_appeared_period'] >= 0:
                            omission = period_index - stat['_last_appeared_period'] - 1
                            
                            # 更新最大遗漏
                            if omission > stat['max_omission']:
                                stat['max_omission'] = omission
                        
                        # 5. 更新最后出现期数
                        stat['_last_appeared_period'] = period_index
                        
                    else:
                        # ===== 号码未出现 =====
                        
                        # 重置连出计数
                        stat['_current_consecutive'] = 0
        
        logger.info("数据扫描完成，开始最终计算")
        
        # 最终计算（基于扫描结果）
        latest_period_index = total_periods - 1
        
        for zone in stats:
            for number in stats[zone]:
                stat = stats[zone][number]
                
                # 1. 计算当前遗漏
                if stat['_last_appeared_period'] >= 0:
                    stat['current_omission'] = latest_period_index - stat['_last_appeared_period']
                else:
                    # 从未出现过，遗漏期数 = 总期数
                    stat['current_omission'] = total_periods
                
                # 2. 计算上期遗漏
                if total_periods >= 2 and stat['_last_appeared_period'] >= 0:
                    stat['last_omission'] = (latest_period_index - 1) - stat['_last_appeared_period']
                    # 如果最后一期刚出现，上期遗漏为0
                    if stat['last_omission'] < 0:
                        stat['last_omission'] = 0
                else:
                    stat['last_omission'] = stat['current_omission']
                
                # 3. 计算平均遗漏
                if stat['occurrence_count'] > 0:
                    stat['avg_omission'] = round(total_periods / stat['occurrence_count'], 2)
                else:
                    stat['avg_omission'] = 0.0
                
                # 4. 更新最大遗漏（考虑当前遗漏）
                if stat['current_omission'] > stat['max_omission']:
                    stat['max_omission'] = stat['current_omission']
                
                # 清理辅助字段
                del stat['_current_consecutive']
                del stat['_last_appeared_period']
        
        logger.info(f"统计计算完成")
        return stats
    
    def save_statistics(self, stats: Dict) -> bool:
        """保存统计数据到数据库（使用批量插入优化性能）"""
        try:
            # 创建统计表（如果不存在）
            self._create_statistics_table()
            
            # 清空当前彩种的旧数据
            logger.info(f"清空 {self.lottery_type} 旧统计数据")
            delete_sql = "DELETE FROM lottery_number_statistics WHERE lottery_type = %s"
            db_manager.execute_query(delete_sql, (self.lottery_type,))
            
            # 准备批量插入数据
            insert_sql = """
            INSERT INTO lottery_number_statistics (
                lottery_type, number_zone, number, 
                occurrence_count, current_omission, last_omission,
                max_omission, avg_omission, max_consecutive,
                updated_at
            ) VALUES (
                %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW()
            )
            """
            
            records = []
            for zone, numbers in stats.items():
                for number, stat in numbers.items():
                    records.append((
                        self.lottery_type,
                        zone,
                        stat['number'],
                        stat['occurrence_count'],
                        stat['current_omission'],
                        stat['last_omission'],
                        stat['max_omission'],
                        stat['avg_omission'],
                        stat['max_consecutive']
                    ))
            
            # 批量插入（使用原生连接以提高性能）
            if records:
                logger.info(f"批量插入 {len(records)} 条统计数据")
                connection = db_manager.mysql_engine.raw_connection()
                try:
                    cursor = connection.cursor()
                    cursor.executemany(insert_sql, records)
                    connection.commit()
                    logger.info(f"✅ 成功保存 {len(records)} 条统计数据")
                    return True
                finally:
                    cursor.close()
                    connection.close()
            
            return False
            
        except Exception as e:
            logger.error(f"❌ 保存统计数据失败: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _create_statistics_table(self):
        """创建统计表（如果不存在）"""
        create_sql = """
        CREATE TABLE IF NOT EXISTS lottery_number_statistics (
            id BIGINT AUTO_INCREMENT PRIMARY KEY,
            lottery_type VARCHAR(10) NOT NULL COMMENT '彩种类型',
            number_zone VARCHAR(10) NOT NULL COMMENT '号码区域(red/blue/all)',
            number INT NOT NULL COMMENT '号码',
            occurrence_count INT NOT NULL DEFAULT 0 COMMENT '出现次数',
            current_omission INT NOT NULL DEFAULT 0 COMMENT '当前遗漏',
            last_omission INT NOT NULL DEFAULT 0 COMMENT '上期遗漏',
            max_omission INT NOT NULL DEFAULT 0 COMMENT '最大遗漏',
            avg_omission DECIMAL(10,2) NOT NULL DEFAULT 0 COMMENT '平均遗漏',
            max_consecutive INT NOT NULL DEFAULT 0 COMMENT '最大连出',
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            UNIQUE KEY uk_lottery_number (lottery_type, number_zone, number),
            INDEX idx_lottery_type (lottery_type),
            INDEX idx_number_zone (number_zone)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='号码统计表';
        """
        
        connection = db_manager.mysql_engine.raw_connection()
        try:
            cursor = connection.cursor()
            cursor.execute(create_sql)
            connection.commit()
            logger.info("✅ 统计表创建成功或已存在")
        except Exception as e:
            logger.error(f"❌ 创建统计表失败: {e}")
            raise
        finally:
            cursor.close()
            connection.close()
    
    def run(self) -> bool:
        """执行统计分析"""
        logger.info(f"{'='*60}")
        logger.info(f"开始分析 {self.lottery_name} 号码统计")
        logger.info(f"{'='*60}")
        
        # 获取历史数据
        historical_data = self.fetch_historical_data()
        if not historical_data:
            logger.error("❌ 没有历史数据")
            return False
        
        # 计算统计数据
        stats = self.calculate_statistics(historical_data)
        if not stats:
            logger.error("❌ 统计计算失败")
            return False
        
        # 保存统计数据
        success = self.save_statistics(stats)
        
        if success:
            logger.info(f"✅ {self.lottery_name} 号码统计分析完成")
            self._print_summary(stats)
        
        return success
    
    def _print_summary(self, stats: Dict):
        """打印统计摘要"""
        print(f"\n{'='*80}")
        print(f"📊 {self.lottery_name} 号码统计摘要")
        print(f"{'='*80}")
        
        for zone, numbers in stats.items():
            print(f"\n【{zone.upper()} 区】共 {len(numbers)} 个号码")
            print(f"{'号码':<6} {'出现次数':<10} {'当前遗漏':<10} {'最大遗漏':<10} {'平均遗漏':<10} {'最大连出':<10}")
            print("-" * 80)
            
            # 按号码排序
            sorted_numbers = sorted(numbers.items(), key=lambda x: x[0])
            
            # 显示所有号码或前20个
            display_count = min(20, len(sorted_numbers))
            for number, stat in sorted_numbers[:display_count]:
                print(f"{stat['number']:<6} {stat['occurrence_count']:<10} {stat['current_omission']:<10} "
                      f"{stat['max_omission']:<10} {stat['avg_omission']:<10.2f} {stat['max_consecutive']:<10}")
            
            if len(sorted_numbers) > display_count:
                print(f"... 还有 {len(sorted_numbers) - display_count} 个号码")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='彩票号码统计分析工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例：
  python number_statistics_calculator.py --lottery ssq
  python number_statistics_calculator.py --lottery dlt --recalculate
  python number_statistics_calculator.py --all
        """
    )
    parser.add_argument('--lottery', 
                       choices=['ssq', 'dlt', 'fc3d', 'pl3', 'pl5', 'qlc', 'qxc', 'kl8'],
                       help='彩种类型')
    parser.add_argument('--all', action='store_true',
                       help='计算所有已启用的彩种')
    parser.add_argument('--recalculate', action='store_true',
                       help='强制重新计算（清空旧数据）')
    
    args = parser.parse_args()
    
    if not args.lottery and not args.all:
        parser.error("请指定 --lottery 或 --all 参数")
    
    # 配置日志
    logger.add(
        "logs/statistics_{time:YYYY-MM-DD}.log",
        rotation="1 day",
        retention="30 days",
        level="INFO"
    )
    
    # 确定要计算的彩种列表
    if args.all:
        lotteries = [lt for lt, config in LOTTERY_TYPES.items() if config.get('enabled', False)]
    else:
        lotteries = [args.lottery]
    
    # 执行统计
    success_count = 0
    for lottery in lotteries:
        calculator = NumberStatisticsCalculator(lottery)
        if calculator.run():
            success_count += 1
        print()  # 空行分隔
    
    # 输出总结
    print(f"\n{'='*80}")
    if success_count == len(lotteries):
        print(f"✅ 全部完成！成功计算 {success_count}/{len(lotteries)} 个彩种")
        print(f"💾 数据已保存到 lottery_number_statistics 表")
    else:
        print(f"⚠️  部分完成：成功 {success_count}/{len(lotteries)} 个彩种")
    print(f"{'='*80}")


if __name__ == '__main__':
    main()
